# 新增颜色识别相关的导入
import cv2
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import os

def detect_color_areas(image_path):
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        # 尝试使用不同的路径分隔符
        image_path_alt = image_path.replace('\\', '/')
        image = cv2.imread(image_path_alt)
        if image is None:
            # 尝试使用双反斜杠
            image_path_double_backslash = image_path_alt.replace('/', '\\')
            image = cv2.imread(image_path_double_backslash)
            if image is None:
                raise FileNotFoundError(f"无法找到文件：{image_path}")
    # 转换为RGB颜色空间
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 定义颜色范围
    # 粉色 #fcc0da 转换为 RGB
    pink_lower = np.array([252, 192, 218])
    pink_upper = np.array([252, 192, 218])
    # 红色 #f7376a 转换为 RGB
    red_lower = np.array([247, 55, 106])
    red_upper = np.array([247, 55, 106])
    # 绿色（考虑到绿白相间的情况）
    green_lower = np.array([0, 100, 0])
    green_upper = np.array([100, 255, 100])
    
    # 创建颜色掩膜
    pink_mask = cv2.inRange(image_rgb, pink_lower, pink_upper)
    red_mask = cv2.inRange(image_rgb, red_lower, red_upper)
    green_mask = cv2.inRange(image_rgb, green_lower, green_upper)
    
    # 打印调试信息
    print(f"粉色掩膜中的非零像素数: {np.count_nonzero(pink_mask)}")
    print(f"红色掩膜中的非零像素数: {np.count_nonzero(red_mask)}")
    print(f"绿色掩膜中的非零像素数: {np.count_nonzero(green_mask)}")
    
    # 寻找轮廓
    pink_contours, _ = cv2.findContours(pink_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    red_contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    green_contours, _ = cv2.findContours(green_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 输出检测到的各颜色区域数量
    print(f"检测到的粉色区域数量: {len(pink_contours)}")
    print(f"检测到的红色区域数量: {len(red_contours)}")
    print(f"检测到的绿色区域数量: {len(green_contours)}")
    
    # 绘制轮廓
    cv2.drawContours(image_rgb, pink_contours, -1, (255, 0, 0), 2)  # 粉色区域用蓝色轮廓表示
    cv2.drawContours(image_rgb, red_contours, -1, (0, 255, 0), 2)   # 红色区域用绿色轮廓表示
    cv2.drawContours(image_rgb, green_contours, -1, (0, 0, 255), 2)  # 绿色区域用红色轮廓表示
    
    # 显示结果图像
    plt.imshow(image_rgb)
    plt.title('Detected Color Areas')
    plt.show()
    
    return pink_contours, red_contours, green_contours

def find_closest_pink_to_green(pink_contours, green_contours, red_contours):
    # 创建一个距离矩阵来存储每个绿色区域到粉色区域的距离
    distance_matrix = np.zeros((len(green_contours), len(pink_contours)))
    
    # 计算每个绿色区域与所有粉色区域之间的距离
    for i, green_cnt in enumerate(green_contours):
        green_moments = cv2.moments(green_cnt)
        if green_moments["m00"] == 0:
            green_center = (0, 0)
        else:
            green_center = (int(green_moments["m10"] / green_moments["m00"]), 
                             int(green_moments["m01"] / green_moments["m00"]))
        
        for j, pink_cnt in enumerate(pink_contours):
            pink_moments = cv2.moments(pink_cnt)
            if pink_moments["m00"] == 0:
                pink_center = (0, 0)
            else:
                pink_center = (int(pink_moments["m10"] / pink_moments["m00"]), 
                                 int(pink_moments["m01"] / pink_moments["m00"]))
            
            # 计算欧氏距离
            distance = np.sqrt(np.sum((np.array(green_center) - np.array(pink_center))**2))
            distance_matrix[i, j] = distance
    
    # 找到每个绿色区域最近的粉色区域
    closest_pink_indices = np.argmin(distance_matrix, axis=1)
    
    # 存储结果
    results = []
    for i, idx in enumerate(closest_pink_indices):
        results.append({
            'green_index': i,
            'closest_pink_index': idx,
            'distance': distance_matrix[i, idx],
            'pink_area': cv2.contourArea(pink_contours[idx])
        })
    
    # 根据粉色区域的大小排序，找到粉色最多的区域
    sorted_results = sorted(results, key=lambda x: x['pink_area'], reverse=True)
    
    
    return sorted_results

def main():
    # 获取当前脚本所在的目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 构建图像文件的绝对路径
    image_path = os.path.join(current_dir, '1.jpg')
    
    # 检测颜色区域
    pink_contours, red_contours, green_contours = detect_color_areas(image_path)
    
    if not pink_contours:
        print("未检测到粉色区域。")
    if not red_contours:
        print("未检测到红色区域。")
    if not green_contours:
        print("未检测到绿色区域。")
    
    # 找到每个绿色数字标识地块附近距离最近且粉色最多的区域
    try:
        sorted_results = find_closest_pink_to_green(pink_contours, green_contours, red_contours)
        
        if sorted_results:
            # 输出最终结果
            print("\n最终结果：")
            for result in sorted_results:
                print(f"绿色区域 {result['green_index']} 最近的粉色区域索引：{result['closest_pink_index']}")
                print(f"距离：{result['distance']}")
                print(f"粉色区域面积：{result['pink_area']}")
                print('---')
        else:
            print("没有找到符合条件的结果。")
    except Exception as e:
        print(f"处理过程中发生错误：{e}")
    
if __name__ == "__main__":
    main()
